#include <fstream>
#include <iostream>
#include <string>
#include <sstream>
#include <array>
#include <vector>
#include <chrono>  //  高精度计时库

#define GLM_FORCE_RADIANS   // 强制glm库使用弧度
#define GLM_FORCE_DEPTH_ZERO_TO_ONE  //  由于vulkan与opengl深度范围不一样，它是0~1.0的
#include "glm/glm.hpp"
#include "glm/gtc/matrix_transform.hpp"

//#include "vulkan/vulkan.h"
//#include "kompute/Core.hpp"

//#include "fmt/src/format.cc"  // 定义了 FMT_HEADER_ONLY 才需要
#include "spdlog/spdlog.h"
#include "spdlog_custom.h"
#include "kpextSpvData.h"
#include "kompute/Tensor.hpp"
#include "kompute/Manager.hpp"
#include "kompute/Algorithm.hpp"
#include "kompute/Sequence.hpp"
#include "kompute/operations/OpTensorSyncLocal.hpp"
#include "kompute/operations/OpTensorSyncDevice.hpp"

using namespace std;

static const unsigned int spirv_demo[] =
{
 0x07230203, 0x00010000, 0x0008000a, 0x00000042, 0x00000000, 0x00020011, 0x00000001, 0x0006000b,
0x00000001, 0x4c534c47, 0x6474732e, 0x3035342e, 0x00000000, 0x0003000e, 0x00000000, 0x00000001,
0x0006000f, 0x00000005, 0x00000004, 0x6e69616d, 0x00000000, 0x0000000b, 0x00060010, 0x00000004,
0x00000011, 0x00000001, 0x00000001, 0x00000001, 0x00030003, 0x00000002, 0x000001c2, 0x00040005,
0x00000004, 0x6e69616d, 0x00000000, 0x00040005, 0x00000008, 0x65646e69, 0x00000078, 0x00080005,
0x0000000b, 0x475f6c67, 0x61626f6c, 0x766e496c, 0x7461636f, 0x496e6f69, 0x00000044, 0x00050005,
0x00000011, 0x5f667562, 0x5f74756f, 0x00000061, 0x00050006, 0x00000011, 0x00000000, 0x5f74756f,
0x00000061, 0x00030005, 0x00000013, 0x00000000, 0x00050005, 0x00000019, 0x5f667562, 0x615f6e69,
0x00000000, 0x00050006, 0x00000019, 0x00000000, 0x615f6e69, 0x00000000, 0x00030005, 0x0000001b,
0x00000000, 0x00050005, 0x00000021, 0x5f667562, 0x625f6e69, 0x00000000, 0x00050006, 0x00000021,
0x00000000, 0x625f6e69, 0x00000000, 0x00030005, 0x00000023, 0x00000000, 0x00050005, 0x0000002f,
0x5f667562, 0x5f74756f, 0x00000062, 0x00050006, 0x0000002f, 0x00000000, 0x5f74756f, 0x00000062,
0x00030005, 0x00000031, 0x00000000, 0x00050005, 0x00000033, 0x736e6f63, 0x6e6f5f74, 0x00000065,
0x00060005, 0x00000034, 0x68737550, 0x736e6f43, 0x746e6174, 0x00000073, 0x00040006, 0x00000034,
0x00000000, 0x006c6176, 0x00050005, 0x00000036, 0x68737570, 0x6e6f635f, 0x00007473, 0x00040047,
0x0000000b, 0x0000000b, 0x0000001c, 0x00040047, 0x00000010, 0x00000006, 0x00000004, 0x00050048,
0x00000011, 0x00000000, 0x00000023, 0x00000000, 0x00030047, 0x00000011, 0x00000003, 0x00040047,
0x00000013, 0x00000022, 0x00000000, 0x00040047, 0x00000013, 0x00000021, 0x00000002, 0x00040047,
0x00000018, 0x00000006, 0x00000004, 0x00050048, 0x00000019, 0x00000000, 0x00000023, 0x00000000,
0x00030047, 0x00000019, 0x00000003, 0x00040047, 0x0000001b, 0x00000022, 0x00000000, 0x00040047,
0x0000001b, 0x00000021, 0x00000000, 0x00040047, 0x00000020, 0x00000006, 0x00000004, 0x00050048,
0x00000021, 0x00000000, 0x00000023, 0x00000000, 0x00030047, 0x00000021, 0x00000003, 0x00040047,
0x00000023, 0x00000022, 0x00000000, 0x00040047, 0x00000023, 0x00000021, 0x00000001, 0x00040047,
0x0000002e, 0x00000006, 0x00000004, 0x00050048, 0x0000002f, 0x00000000, 0x00000023, 0x00000000,
0x00030047, 0x0000002f, 0x00000003, 0x00040047, 0x00000031, 0x00000022, 0x00000000, 0x00040047,
0x00000031, 0x00000021, 0x00000003, 0x00040047, 0x00000033, 0x00000001, 0x00000000, 0x00050048,
0x00000034, 0x00000000, 0x00000023, 0x00000000, 0x00030047, 0x00000034, 0x00000002, 0x00040047,
0x00000041, 0x0000000b, 0x00000019, 0x00020013, 0x00000002, 0x00030021, 0x00000003, 0x00000002,
0x00040015, 0x00000006, 0x00000020, 0x00000000, 0x00040020, 0x00000007, 0x00000007, 0x00000006,
0x00040017, 0x00000009, 0x00000006, 0x00000003, 0x00040020, 0x0000000a, 0x00000001, 0x00000009,
0x0004003b, 0x0000000a, 0x0000000b, 0x00000001, 0x0004002b, 0x00000006, 0x0000000c, 0x00000000,
0x00040020, 0x0000000d, 0x00000001, 0x00000006, 0x0003001d, 0x00000010, 0x00000006, 0x0003001e,
0x00000011, 0x00000010, 0x00040020, 0x00000012, 0x00000002, 0x00000011, 0x0004003b, 0x00000012,
0x00000013, 0x00000002, 0x00040015, 0x00000014, 0x00000020, 0x00000001, 0x0004002b, 0x00000014,
0x00000015, 0x00000000, 0x00030016, 0x00000017, 0x00000020, 0x0003001d, 0x00000018, 0x00000017,
0x0003001e, 0x00000019, 0x00000018, 0x00040020, 0x0000001a, 0x00000002, 0x00000019, 0x0004003b,
0x0000001a, 0x0000001b, 0x00000002, 0x00040020, 0x0000001d, 0x00000002, 0x00000017, 0x0003001d,
0x00000020, 0x00000017, 0x0003001e, 0x00000021, 0x00000020, 0x00040020, 0x00000022, 0x00000002,
0x00000021, 0x0004003b, 0x00000022, 0x00000023, 0x00000002, 0x00040020, 0x00000029, 0x00000002,
0x00000006, 0x0003001d, 0x0000002e, 0x00000006, 0x0003001e, 0x0000002f, 0x0000002e, 0x00040020,
0x00000030, 0x00000002, 0x0000002f, 0x0004003b, 0x00000030, 0x00000031, 0x00000002, 0x00040032,
0x00000017, 0x00000033, 0x00000000, 0x0003001e, 0x00000034, 0x00000017, 0x00040020, 0x00000035,
0x00000009, 0x00000034, 0x0004003b, 0x00000035, 0x00000036, 0x00000009, 0x00040020, 0x00000037,
0x00000009, 0x00000017, 0x0004002b, 0x00000006, 0x00000040, 0x00000001, 0x0006002c, 0x00000009,
0x00000041, 0x00000040, 0x00000040, 0x00000040, 0x00050036, 0x00000002, 0x00000004, 0x00000000,
0x00000003, 0x000200f8, 0x00000005, 0x0004003b, 0x00000007, 0x00000008, 0x00000007, 0x00050041,
0x0000000d, 0x0000000e, 0x0000000b, 0x0000000c, 0x0004003d, 0x00000006, 0x0000000f, 0x0000000e,
0x0003003e, 0x00000008, 0x0000000f, 0x0004003d, 0x00000006, 0x00000016, 0x00000008, 0x0004003d,
0x00000006, 0x0000001c, 0x00000008, 0x00060041, 0x0000001d, 0x0000001e, 0x0000001b, 0x00000015,
0x0000001c, 0x0004003d, 0x00000017, 0x0000001f, 0x0000001e, 0x0004003d, 0x00000006, 0x00000024,
0x00000008, 0x00060041, 0x0000001d, 0x00000025, 0x00000023, 0x00000015, 0x00000024, 0x0004003d,
0x00000017, 0x00000026, 0x00000025, 0x00050085, 0x00000017, 0x00000027, 0x0000001f, 0x00000026,
0x0004006d, 0x00000006, 0x00000028, 0x00000027, 0x00060041, 0x00000029, 0x0000002a, 0x00000013,
0x00000015, 0x00000016, 0x0004003d, 0x00000006, 0x0000002b, 0x0000002a, 0x00050080, 0x00000006,
0x0000002c, 0x0000002b, 0x00000028, 0x00060041, 0x00000029, 0x0000002d, 0x00000013, 0x00000015,
0x00000016, 0x0003003e, 0x0000002d, 0x0000002c, 0x0004003d, 0x00000006, 0x00000032, 0x00000008,
0x00050041, 0x00000037, 0x00000038, 0x00000036, 0x00000015, 0x0004003d, 0x00000017, 0x00000039,
0x00000038, 0x00050085, 0x00000017, 0x0000003a, 0x00000033, 0x00000039, 0x0004006d, 0x00000006,
0x0000003b, 0x0000003a, 0x00060041, 0x00000029, 0x0000003c, 0x00000031, 0x00000015, 0x00000032,
0x0004003d, 0x00000006, 0x0000003d, 0x0000003c, 0x00050080, 0x00000006, 0x0000003e, 0x0000003d,
0x0000003b, 0x00060041, 0x00000029, 0x0000003f, 0x00000031, 0x00000015, 0x00000032, 0x0003003e,
0x0000003f, 0x0000003e, 0x000100fd, 0x00010038,  };

static std::vector<uint32_t> compileSource(const std::string& source)
{
    std::ofstream fileOut("tmp_kp_shader.comp");
    fileOut << source;
    fileOut.close();
    if (system(std::string("glslangValidator -V tmp_kp_shader.comp -o tmp_kp_shader.comp.spv").c_str()))
        throw std::runtime_error("Error running glslangValidator command");
    std::ifstream fileStream("tmp_kp_shader.comp.spv", std::ios::binary);
//    std::vector<char> buffer;
    //    std::istreambuf_iterator<T>(in) 使用流生成访问迭代器，流访问迭代器到达未尾相当于 istreambuf_iterator<T>()
//    buffer.insert(buffer.begin(), std::istreambuf_iterator<char>(fileStream), istreambuf_iterator<char>());
    std::vector<char> buffer((std::istreambuf_iterator<char>(fileStream)),
                             std::istreambuf_iterator<char>()
                             );
    std::vector<unsigned char> &buf1 = * reinterpret_cast<std::vector<unsigned char> *>(&buffer);
    auto str_spirv = fmt::format("{:#04x}", fmt::join(buf1, ", "));  // "{:#04x}" 包括0x前缀总共4位
    fmt::print("spirv : \n {} \n {:*<30} \n", str_spirv, "*");

    std::vector<uint32_t> ret{reinterpret_cast<uint32_t*>(buffer.data()),
                reinterpret_cast<uint32_t*>(buffer.data() + buffer.size())};
    stringstream ss;
    int i = 0;
    for(auto & item : ret){
        ss << fmt::format("{:#010x}, ", item);
        i++;
        if(i == 8) ss << std::endl;
        i = i % 8;
    }
    string s2 = ss.str();
    fmt::print("static const unsigned int spirv_demo[] = \n{{ \n {} }}; \n {:*<30} \n", s2, "*");
    return ret;
}

void kompute(const std::string& shader) {

    // 1. Create Kompute Manager with default settings (device 0, first queue and no extensions)
    kp::Manager mgr;

    // 2. Create and initialise Kompute Tensors through manager

    // Default tensor constructor simplifies creation of float values
    auto tensorInA = mgr.tensor({ 2., 2., 2. });
    auto tensorInB = mgr.tensor({ 1., 2., 3. });
    // Explicit type constructor supports uint32, int32, double, float and bool
    auto tensorOutA = mgr.tensorT<uint32_t>({ 0, 0, 0 });
    auto tensorOutB = mgr.tensorT<uint32_t>({ 0, 0, 0 });

    std::vector<std::shared_ptr<kp::Tensor>> params = {tensorInA, tensorInB, tensorOutA, tensorOutB};

    // 3. Create algorithm based on shader (supports buffers & push/spec constants)
    kp::Workgroup workgroup{3, 1, 1};  // c++11 std::array 不能用括号加{}初始化
    std::vector<float> specConsts({ 2 });
    std::vector<float> pushConstsA({ 2.0 });
    std::vector<float> pushConstsB({ 3.0 });

    std::vector<uint32_t> orgspirv = compileSource(shader);
//    std::vector<uint32_t> tspriv;
//    tspriv.insert(tspriv.begin(),
//                  reinterpret_cast<const uint32_t *>(spirv_demo),
//                  reinterpret_cast<const uint32_t *>(spirv_demo + sizeof (spirv_demo)) + 1
//                  );

    std::vector<uint32_t> tspriv(spirv_demo,
                                 spirv_demo + sizeof (spirv_demo) / sizeof (spirv_demo[0]) //sizeof (decltype(spirv_demo[0]))
                                 );
//    std::vector<uint32_t> tspriv(ExtendAndWin,
//                                 ExtendAndWin + sizeof (ExtendAndWin) / sizeof (ExtendAndWin[0]) //sizeof (decltype(spirv_demo[0]))
//                                 );
    KP_LOG_INFO("is compiled equel to showed : {}", orgspirv == tspriv);
    auto algorithm = mgr.algorithm(params,
                                   // See documentation shader section for compileSource
                                   tspriv, //compileSource(shader),
                                   workgroup,
                                   specConsts,
                                   pushConstsA);

    // 4. Run operation synchronously using sequence
    mgr.sequence()
        ->record<kp::OpTensorSyncDevice>(params)
        ->record<kp::OpAlgoDispatch>(algorithm) // Binds default push consts
        ->eval() // Evaluates the two recorded operations
        ->record<kp::OpAlgoDispatch>(algorithm, pushConstsB) // Overrides push consts
        ->eval(); // Evaluates only last recorded operation

    // 5. Sync results from the GPU asynchronously
    auto sq = mgr.sequence();
    sq->evalAsync<kp::OpTensorSyncLocal>(params);

    // ... Do other work asynchronously whilst GPU finishes

    sq->evalAwait();

//    fmt::print("{0:*^{1}}\n","result",30);
    KP_LOG_INFO("{0:*^{1}}","result",30);
    KP_LOG_INFO(fmt::format("tensorOutA: {}", fmt::join(tensorOutA->vector(), ", ")));
    KP_LOG_INFO(fmt::format("tensorOutB: {}", fmt::join(tensorOutB->vector(), ", ")));
//    // Prints the first output which is: { 4, 8, 12 }
//    for (const float& elem : tensorOutA->vector()) std::cout << elem << "  ";
//    std::cout << std::endl;
//    // Prints the second output which is: { 10, 10, 10 }
//    for (const float& elem : tensorOutB->vector()) std::cout << elem << "  ";
//    std::cout << std::endl;

} // Manages / releases all CPU and GPU memory resources


int main()
{
//    kpextSetLogger("kpextLogger", 1, 2);
//    kpextSetLogger(nullptr, 1, 2);

#ifdef DEBUG
    spdlog::set_level(spdlog::level::debug);
    SPDLOG_LOGGER_DEBUG(spdlog::default_logger_raw(), "spdlog::default_logger_raw");
    spdlog::set_level(spdlog::level::info);
    KP_LOG_DEBUG("this debug log will not show \n");
    KP_LOG_INFO("this info log will show","\n");
#endif
//    spdlog::set_level(spdlog::level::off);

    // Define a raw string shader (or use the Kompute tools to compile to SPIRV / C++ header
    // files). This shader shows some of the main components including constants, buffers, etc
    std::string shader = (R"(
            #version 450

            layout (local_size_x = 1) in;

            // The input tensors bind index is relative to index in parameter passed
            layout(set = 0, binding = 0) buffer buf_in_a { float in_a[]; };
            layout(set = 0, binding = 1) buffer buf_in_b { float in_b[]; };
            layout(set = 0, binding = 2) buffer buf_out_a { uint out_a[]; };
            layout(set = 0, binding = 3) buffer buf_out_b { uint out_b[]; };

            // Kompute supports push constants updated on dispatch
            layout(push_constant) uniform PushConstants {
                float val;
            } push_const;

            // Kompute also supports spec constants on initalization
            layout(constant_id = 0) const float const_one = 0;

            void main() {
                uint index = gl_GlobalInvocationID.x;
                out_a[index] += uint( in_a[index] * in_b[index] );
                out_b[index] += uint( const_one * push_const.val );
            }
        )");

    kompute(shader);

    return 0;
}
